import os
import cv2
import numpy as np
import matplotlib.pyplot as plt

TEMPORARY = os.path.join(os.path.dirname(__file__), 'Temporary.png')
CMAPS = {
    '235U': 'spring',
    '238U': 'summer'
}
plt.rcParams['font.sans-serif'] = ['SimHei']

class Plot:
    def __init__(self, data):
        self.CHAIN_NAMES = ['235U', '238U']
        self.T = data['t']
        self.METHOD = data['Method']
        if self.METHOD == 'Power':
            self.POWER = data['Power']
        self.DATA = data['Data']
        self.NUCLIDE_NUM = {}
        self.COLORS = {}
        for name in self.CHAIN_NAMES:
            self.NUCLIDE_NUM[name] = len(self.DATA[name])
            self.COLORS[name] = plt.cm.get_cmap(name=CMAPS[name], lut=2 * self.NUCLIDE_NUM[name])
        self.FIG, self.AX = plt.subplots()

        self.curves = {
            '235U': {},
            '238U': {}
        }
        self.N_MAX = {
            '235U': {},
            '238U': {}
        }

        self.result = None

    def Draw(self, draw_list=None):
        if draw_list is None:
            draw_list = {}
            for name in self.CHAIN_NAMES:
                draw_list[name] = {}
                for nuclide in self.DATA[name].keys():
                    draw_list[name][nuclide] = {
                        'Enable': True,
                        'Ratio': 1
                    }

        if self.METHOD == 'Flux':
            time_sequence = np.array(self.DATA['t']) / (24 * 60 * 60)
            if time_sequence[-1] > 3650:
                time_sequence /= 365
                T_UNIT = '年'
            elif time_sequence[-1] > 300:
                time_sequence /= 30
                T_UNIT = '月'
            else:
                T_UNIT = '天'
            plt.title("同位素成分随运行时间变化")
            plt.xlabel("运行时间/(%s)"%T_UNIT)
        else:
            time_sequence = np.array(self.DATA['t']) / (24 * 60 * 60) * self.POWER * 1e-3
            plt.title("同位素成分随燃耗深度变化")
            plt.xlabel("燃耗深度/($GW\cdot d\cdot {t}^{-1}$)")
        plt.ylabel("同位素成分/($kg\cdot{t}^{-1}$)")

        nuclide_id = -1
        for name in self.CHAIN_NAMES:
            for nuclide in draw_list[name]:
                nuclide_id += 1
                if nuclide in self.DATA[name].keys() and draw_list[name][nuclide]['Enable'] and (nuclide not in self.curves[name].keys()):
                    self.Add_Curve(name, nuclide, nuclide_id, draw_list, time_sequence)
                elif (not draw_list[name][nuclide]['Enable']) and nuclide in self.curves[name].keys():
                    self.Remove_Curve(name, nuclide)
                elif draw_list[name][nuclide]['Enable'] and nuclide in self.curves[name].keys() and draw_list[name][nuclide]['Ratio'] != self.N_MAX[name][nuclide]['Ratio']:
                    self.Remove_Curve(name, nuclide)
                    self.Add_Curve(name, nuclide, nuclide_id, draw_list, time_sequence)

        N_max = self.Find_N_max()
        plt.ylim((-0.05 * N_max, 1.05 * N_max))

        plt.grid()
        plt.legend()
        plt.tight_layout()

        plt.savefig(TEMPORARY)
        self.result = cv2.imread(TEMPORARY, cv2.IMREAD_UNCHANGED)
        os.remove(TEMPORARY)

    def Add_Curve(self, name, nuclide, nuclide_id, draw_list, time_sequence):
        N_sequence = np.array(self.DATA[name][nuclide]) * draw_list[name][nuclide]['Ratio'] * int(nuclide[0:3]) * 1e-3
        nuclide_label = nuclide if draw_list[name][nuclide]['Ratio'] == 1 else nuclide + ' * %f'%draw_list[name][nuclide]['Ratio']
        self.curves[name][nuclide], = plt.plot(time_sequence, N_sequence, color=self.COLORS[name](nuclide_id), label=nuclide_label)
        # “,”很关键
        self.N_MAX[name][nuclide] = {
            "Value": np.max(N_sequence),
            "Ratio": draw_list[name][nuclide]['Ratio']
        }

    def Remove_Curve(self, name, nuclide):
        self.curves[name][nuclide].remove()
        del self.curves[name][nuclide]
        del self.N_MAX[name][nuclide]

    def Find_N_max(self):
        N_maxs = []
        for name in self.CHAIN_NAMES:
            for iterm in self.N_MAX[name].values():
                N_maxs.append(iterm['Value'])
        return np.max(N_maxs)

    def Show(self, delay=0):
        if self.result is not None:
            cv2.imshow('Result', self.result)
            cv2.waitKey(delay)

    def Save(self, *args, **kwargs):
        plt.savefig(*args, **kwargs)